package wang.yang;

import javafx.util.Pair;
import org.apache.commons.collections.CollectionUtils;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Random;

/**
 * 竞选主节点
 * 竞选成功
 * 找出down掉的worker，重新分配任务
 */
public class MyMaster {
    private final static Logger LOG = LoggerFactory.getLogger(MyMaster.class);
    private static final byte[] ZERO_LENGTH_ARR = new byte[0];

    ZooKeeper zk;
    String hostPort;
    String serverId = Integer.toHexString(new Random().nextInt());
    private boolean connected = false;
    private boolean expired = false;

    enum MasterStates {RUNNING, ELECTED, NOTELECTED}

    ;
    private List<String> unassignedTasksCache;
    private List<String> activeWorkersCache;

    private volatile MasterStates state = MasterStates.RUNNING;

    MasterStates getState() {
        return state;
    }

    MyMaster(String hostPort) {
        this.hostPort = hostPort;
    }

    void startZK() {
        try {
            zk = new ZooKeeper(hostPort, 5000, event -> {
                LOG.info("{}收到服务器事件{}", serverId, event);
                if (Watcher.Event.EventType.None == event.getType()) {
                    switch (event.getState()) {
                        case SyncConnected:
                            connected = true;
                            break;
                        case Disconnected:
                            connected = false;
                            break;
                        case Expired:
                            expired = true;
                            connected = false;
                            LOG.error("{} session expired", serverId);
                        default:
                            break;
                    }
                }
            });
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    void stopZK() throws Exception {
        zk.close();
    }

    private void masterExists() {
        zk.exists("/master", (event) -> {
            if (event.getType() == Watcher.Event.EventType.NodeDeleted) {
                // 监听到主节点被删除，参与竞选
                if ("/master".equals(event.getPath())) {
                    runForMaster();
                }
            }
        }, masterExistsCallback, null);
    }

    AsyncCallback.StringCallback masterCreateCallback = new AsyncCallback.StringCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, String name) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getMasterNode();
                    break;
                case OK:
                    state = MasterStates.ELECTED;
                    takeLeadership();
                    break;
                case NODEEXISTS:
                    state = MasterStates.NOTELECTED;
                    masterExists();
                    break;
                default:
                    state = MasterStates.NOTELECTED;
                    LOG.error("Something went wrong when running for master.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
            LOG.info(serverId + (state == MasterStates.ELECTED ? "已" : "未") + "竞选为主节点");
        }
    };

    AsyncCallback.StatCallback masterExistsCallback = new AsyncCallback.StatCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    masterExists();
                    break;
                case OK:
                    // 监听master节点删除
                    break;
                case NONODE:
                    state = MasterStates.RUNNING;
                    runForMaster();
                    break;
                default:
                    getMasterNode();
                    break;
            }
        }
    };

    public void bootstrap() {
        createParent("/workers", ZERO_LENGTH_ARR);
        createParent("/assign", ZERO_LENGTH_ARR);
        createParent("/unassignedTasks", ZERO_LENGTH_ARR);
        createParent("/status", ZERO_LENGTH_ARR);
    }

    void createParent(String path, byte[] data) {
        zk.create(path,
                data,
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT,
                createParentCallback,
                data);
    }

    AsyncCallback.StringCallback createParentCallback = new AsyncCallback.StringCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, String name) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    createParent(path, (byte[]) ctx);
                    break;
                case OK:
                    LOG.info("Parent created");
                    break;
                case NODEEXISTS:
                    LOG.warn("Parent already registered: " + path);
                    break;
                default:
                    LOG.error("Something went wrong: ",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    void takeLeadership() {
        // 已成为主节点
        LOG.info("Going for list of workers");
        // 构建本地信息，清理不一致数据，同步
        constructMaster();
        // 监听worker变更，有变化则同步
        syncWorkerChange();
        // 分配任务，监听任务变更
        getUnassignedTasks();
    }

    private void syncWorkerChange() {
        zk.getChildren("/workers", (event) -> {
            if (event.getType() == Watcher.Event.EventType.NodeChildrenChanged) {
                syncWorkerChange();
            }
        }, syncWorkerChangeCallback, null);
    }

    AsyncCallback.ChildrenCallback syncWorkerChangeCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> activeWorkers) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getMasterNode();
                    break;
                case OK:
                    List<String> absentWorkers = new ArrayList<>(activeWorkersCache);
                    absentWorkers.removeAll(activeWorkers);
                    activeWorkersCache = activeWorkers;
                    for (String absentWorker : absentWorkers) {
                        // 重新分配任务
                        reassignTasks(absentWorker);
                    }
                    break;
                default:
                    LOG.error("Error when reading data.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    private void reassignTasks(String absentWorker) {
        zk.getChildren("/assign/" + absentWorker, false, getAbsentWorkerTasksCallback, absentWorker);
    }

    AsyncCallback.ChildrenCallback getAbsentWorkerTasksCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object absentWorker, List<String> absentWorkerTasks) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    reassignTasks((String) absentWorker);
                    break;
                case OK:
                    if (CollectionUtils.isNotEmpty(absentWorkerTasks)) {
                        for (String absentWorkerTask : absentWorkerTasks) {
                            String taskPath = "/assign/" + absentWorker + "/" + absentWorkerTask;
                            zk.getData(taskPath, false, getAbsentWorkerTaskDataCallback, taskPath);
                        }
                    }
                    break;
                default:
                    LOG.error("Error when reading data.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    AsyncCallback.DataCallback getAbsentWorkerTaskDataCallback = new AsyncCallback.DataCallback() {
        @Override
        public void processResult(int rc, String path, Object taskPath, byte[] taskData, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    zk.getData((String) taskPath, false, getAbsentWorkerTaskDataCallback, taskPath);
                    break;
                case OK:
                    // 先删除已分配任务节点，再创建新任务
                    Pair<String, String> taskPathAndTaskData = new Pair<>((String) taskPath, new String(taskData));
                    deleteAssignedTaskAndCreateTask(taskPathAndTaskData);
                    break;
                default:
                    LOG.error("Error when reading data.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    private void deleteAssignedTaskAndCreateTask(Pair<String, String> taskPathAndTaskData) {
        String taskPath = taskPathAndTaskData.getKey();
        String taskData = taskPathAndTaskData.getValue();
        String taskId = taskPath.toString().substring(taskPath.toString().lastIndexOf('/') + 1);
        Transaction transaction = zk.transaction();
        transaction.delete(taskPath, -1);
        transaction.create("/unassignedTasks/" + taskId, taskData.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        transaction.commit(deleteAssignedTaskAndCreateTaskCallback, taskPathAndTaskData);
    }

    AsyncCallback.MultiCallback deleteAssignedTaskAndCreateTaskCallback = new AsyncCallback.MultiCallback() {
        @Override
        public void processResult(int rc, String path, Object taskPathAndTaskData, List<OpResult> opResults) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    deleteAssignedTaskAndCreateTask((Pair<String, String>) taskPathAndTaskData);
                    break;
                case OK:
                    OpResult.CreateResult createResult = (OpResult.CreateResult) opResults.get(1);
                    LOG.info("任务再分配成功 {} created", createResult.getPath());
                    break;
                default:
                    LOG.error("Error when reading data.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    private void constructMaster() {
        // 获取activeWorkers
        try {
            activeWorkersCache = zk.getChildren("/workers", false);
            List<String> assignedWorkers = zk.getChildren("/assign", false);

            for (String assignedWorker : assignedWorkers) {
                if (!activeWorkersCache.contains(assignedWorker)) {
                    // 为已宕机的worker重新分配任务
                    List<String> assignedTasks = zk.getChildren("/assign/" + assignedWorker, false);
                    for (String assignedTask : assignedTasks) {
                        // 获取任务信息
                        byte[] taskData = zk.getData("/assign/" + assignedWorker + "/" + assignedTask, false, new Stat());
                        // 从已分配列表中删除，重新创建
                        Transaction transaction = zk.transaction();
                        transaction.delete("/assign/" + assignedWorker + "/" + assignedTask, -1);
                        transaction.create("/unassignedTasks/" + assignedTask, taskData, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                        transaction.commit();
                    }
                }
            }
        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    /*
     * Run for master. To run for master, we try to create the /master znode,
     * with masteCreateCallback being the callback implementation.
     * In the case the create call succeeds, the client becomes the master.
     * If it receives a CONNECTIONLOSS event, then it needs to check if the
     * znode has been created. In the case the znode exists, it needs to check
     * which server is the master.
     */

    /**
     * Tries to create a /master lock znode to acquire leadership.
     */
    public void runForMaster() {
        LOG.info("Running for master");
        // 竞选主节点
        zk.create("/master",
                serverId.getBytes(),
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL,
                masterCreateCallback,
                null);
    }

    AsyncCallback.DataCallback getMasterNodeCallback = new AsyncCallback.DataCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, byte[] masterData, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getMasterNode();
                    break;
                case NONODE:
                    runForMaster();
                    break;
                case OK:
                    if (serverId.equals(new String(masterData))) {
                        state = MasterStates.ELECTED;
                        takeLeadership();
                    } else {
                        state = MasterStates.NOTELECTED;
                        masterExists();
                    }
                    break;
                default:
                    LOG.error("Error when reading data.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    void getMasterNode() {
        // 获取master节点值，确认自己是否为master
        zk.getData("/master", false, getMasterNodeCallback, null);
    }

    /*
     ******************************************************
     ******************************************************
     * Methods for receiving new tasks and assigning them.*
     ******************************************************
     ******************************************************
     */

    Watcher unassignedTasksChangeWatcher = new Watcher() {
        @Override
        public void process(WatchedEvent e) {
            if (e.getType() == Event.EventType.NodeChildrenChanged) {
                getUnassignedTasks();
            }
        }
    };

    void getUnassignedTasks() {
        zk.getChildren("/unassignedTasks",
                unassignedTasksChangeWatcher,
                getUnassignedTasksCallback,
                null);
    }

    AsyncCallback.ChildrenCallback getUnassignedTasksCallback = new AsyncCallback.ChildrenCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, List<String> unassignedTaskList) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getUnassignedTasks();
                    break;
                case OK:
                    List<String> newlyAddedTasks;
                    if (Objects.isNull(unassignedTasksCache)) {
                        unassignedTasksCache = new ArrayList<>(unassignedTaskList);
                        newlyAddedTasks = unassignedTaskList;
                    } else {
                        newlyAddedTasks = new ArrayList<>(unassignedTaskList);
                        newlyAddedTasks.removeAll(unassignedTasksCache);
                        unassignedTasksCache = unassignedTaskList;
                    }
                    if (CollectionUtils.isNotEmpty(newlyAddedTasks)) {
                        assignTasks(newlyAddedTasks);
                    }
                    break;
                default:
                    LOG.error("getChildren failed.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    void assignTasks(List<String> tasks) {
        for (String task : tasks) {
            getTaskData(task);
        }
    }

    void getTaskData(String taskId) {
        zk.getData("/unassignedTasks/" + taskId,
                false,
                getTaskDataCallback,
                taskId);
    }

    AsyncCallback.DataCallback getTaskDataCallback = new AsyncCallback.DataCallback() {
        @Override
        public void processResult(int rc, String path, Object taskId, byte[] data, Stat stat) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getTaskData((String) taskId);
                    break;
                case OK:
                    /*
                     * Choose worker at random.
                     */
                    if (CollectionUtils.isNotEmpty(activeWorkersCache)) {
                        String designatedWorker = activeWorkersCache.get(new Random().nextInt(activeWorkersCache.size()));
                        /*
                         * Assign task to randomly chosen worker.
                         */
                        String assignmentPath = "/assign/" +
                                designatedWorker +
                                "/" +
                                (String) taskId;
                        LOG.info("Assignment path: " + assignmentPath);
                        createAssignment(assignmentPath, data);
                    } else {
                        LOG.info("no active worker");
                    }

                    break;
                default:
                    LOG.error("Error when trying to get task data.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    void createAssignment(String path, byte[] data) {
        zk.create(path,
                data,
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT,
                assignTaskCallback,
                data);
    }

    AsyncCallback.StringCallback assignTaskCallback = new AsyncCallback.StringCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx, String name) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    createAssignment(path, (byte[]) ctx);
                    break;
                case OK:
                    LOG.info("Task assigned correctly: " + name);
                    // 删除任务列表下的任务
                    deleteTask(name.substring(name.lastIndexOf("/") + 1));
                    break;
                case NODEEXISTS:
                    LOG.warn("Task already assigned");
                    break;
                default:
                    LOG.error("Error when trying to assign task.",
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    /*
     * Once assigned, we delete the task from /unassignedtasks
     */
    void deleteTask(String name) {
        zk.delete("/unassignedTasks/" + name, -1, taskDeleteCallback, null);
    }

    AsyncCallback.VoidCallback taskDeleteCallback = new AsyncCallback.VoidCallback() {
        @Override
        public void processResult(int rc, String path, Object ctx) {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    deleteTask(path);
                    break;
                case OK:
                    LOG.info("Successfully deleted " + path);
                    break;
                case NONODE:
                    LOG.info("Task has been deleted already");
                    break;
                default:
                    LOG.error("Something went wrong here, " +
                            KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }
    };

    public static void main(String args[]) throws Exception {
        MyMaster myMaster = new MyMaster("127.0.0.1:2181");
        myMaster.startZK();

        while (!myMaster.connected) {
            Thread.sleep(100);
        }

        myMaster.bootstrap();
        myMaster.runForMaster();

        while (!myMaster.expired) {
            Thread.sleep(100);
        }

        myMaster.stopZK();
    }
}
